from typing import Dict, List, Any, Union
import os

class TreeVisualizer:
    """Visualizes code trees as text."""
    @staticmethod
    def _format_node_name(node_name):
        """Format a node name for better readability by removing duplicate module/class names."""
        # Split the full name into parts
        parts = node_name.split('.')
        
        # If we have at least 3 parts and the first two are identical (like Module.Module.method), 
        # remove the duplicate
        if len(parts) >= 3 and parts[0] == parts[1]:
            return f"{parts[0]}.{'.'.join(parts[2:])}"
        
        return node_name

    @classmethod
    def visualize_call_tree(cls, tree, output_file=None):
        """Visualize a call tree using ASCII art."""
        if not tree:
            print("Empty tree, nothing to visualize")
            return
        
        lines = []
        
        # Sort the root nodes for consistent output
        for i, (root, subtree) in enumerate(sorted(tree.items())):
            # Format the root node name
            formatted_root = cls._format_node_name(root)
            
            # Add the root node
            prefix = "└── " if i == len(tree) - 1 else "├── "
            lines.append(formatted_root)
            
            # Add the children with proper indentation
            cls._add_subtree(subtree, "", i == len(tree) - 1, lines, 1)
        
        # Join all lines
        tree_str = "\n".join(lines)
        
        # Print to console
        print(tree_str)
        
        # Write to file if specified
        if output_file:
            with open(output_file, 'w') as f:
                f.write(tree_str)

    @classmethod
    def _add_subtree(cls, subtree, prefix, is_last, lines, depth):
        """Add a subtree to the visualization with proper indentation."""
        if not subtree:
            return
        
        # Calculate the new prefix for children
        new_prefix = prefix + ("    " if is_last else "│   ")
        
        # Sort the children for consistent output
        sorted_children = sorted(subtree.items())
        
        for i, (node, node_subtree) in enumerate(sorted_children):
            # Format the node name
            formatted_node = cls._format_node_name(node)
            
            # Determine the prefix for this node
            child_prefix = "└── " if i == len(sorted_children) - 1 else "├── "
            
            # Add the node to lines
            lines.append(f"{prefix}{child_prefix}{formatted_node}")
            
            # Add this node's children
            cls._add_subtree(
                node_subtree, 
                new_prefix, 
                i == len(sorted_children) - 1, 
                lines,
                depth + 1
            )

    @staticmethod
    def visualize_call_chains(chains: List[List[str]], output_file: str = None):
        """
        Visualize call chains as text.
        
        Args:
            chains: A list of chains (each chain is a list of function names)
            output_file: Path to output file (if None, print to console)
        """
        if not chains:
            print("Warning: Empty chains provided to visualize_call_chains")
            return ""
            
        output = []
        
        output.append("Call Chains:")
        output.append("="*40)
        
        for i, chain in enumerate(chains):
            output.append(f"Chain {i+1}:")
            for j, func in enumerate(chain):
                if j > 0:
                    output.append("  " * j + "└── " + func)
                else:
                    output.append(func)
            output.append("")
        
        text = "\n".join(output)
        
        if output_file:
            # Ensure output directory exists
            os.makedirs(os.path.dirname(output_file), exist_ok=True)
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(text)
            print(f"Call chains written to {output_file}")
        else:
            print(text)
        
        return text

    @staticmethod
    def visualize_computation_steps(steps, output_file=None):
        """Visualize computation steps for functions."""
        lines = []
        
        for func_name, func_steps in sorted(steps.items()):
            # Format the function name for better readability
            parts = func_name.split('.')
            if len(parts) >= 3 and parts[0] == parts[1]:
                formatted_name = f"{parts[0]}.{'.'.join(parts[2:])}"
            else:
                formatted_name = func_name
            
            lines.append(f"Function: {formatted_name}")
            for i, step in enumerate(func_steps, 1):
                lines.append(f"  {i}. {step}")
            lines.append("")
        
        # Join all lines
        steps_str = "\n".join(lines)
        
        # Print to console
        # print(steps_str)
        
        # Write to file if specified
        if output_file:
            with open(output_file, 'w') as f:
                f.write(steps_str)
